package com.example.vpn_demo.proxy

import android.util.Log
import com.example.vpn_demo.cert.CertificateManager
import com.example.vpn_demo.model.HttpRequest
import com.example.vpn_demo.viewmodel.ViewModelProvider
import java.io.*
import java.net.ServerSocket
import java.net.Socket
import java.util.concurrent.ExecutorService
import java.util.concurrent.Executors
import java.net.InetAddress
import java.net.InetSocketAddress
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.launch
import kotlinx.coroutines.withContext
import kotlinx.coroutines.SupervisorJob
import java.security.KeyStore
import javax.net.ssl.*
import java.util.concurrent.TimeUnit
import okhttp3.OkHttpClient
import okhttp3.Protocol
import okhttp3.Request
import okhttp3.Response
import android.content.Context

class ProxyServer(
    private val certificateManager: CertificateManager,
    private val port: Int = 8888
) {
    companion object {
        private const val TAG = "ProxyServer"
        private const val BUFFER_SIZE = 8192
    }

    private var running = false
    private var serverSocket: ServerSocket? = null
    private val threadPool: ExecutorService = Executors.newCachedThreadPool()
    private val connectionScope = CoroutineScope(Dispatchers.IO + SupervisorJob())

    fun start() {
        if (running) {
            Log.w(TAG, "代理服务器已经在运行")
            return
        }

        running = true
        Log.i(TAG, "开始启动代理服务器...")
        threadPool.execute {
            try {
                serverSocket = ServerSocket().apply {
                    reuseAddress = true
                    bind(InetSocketAddress("127.0.0.1", port))
                }
                Log.i(TAG, "代理服务器成功启动在端口 $port")

                while (running) {
                    try {
                        val socket = serverSocket?.accept() ?: break
                        Log.i(TAG, "收到新连接: ${socket.inetAddress.hostAddress}:${socket.port}")
                        
                        socket.apply {
                            keepAlive = true
                            tcpNoDelay = true
                            soTimeout = 30000
                        }
                        
                        handleConnection(socket)
                    } catch (e: Exception) {
                        if (running) {
                            Log.e(TAG, "接受连接时出错", e)
                        }
                    }
                }
            } catch (e: Exception) {
                Log.e(TAG, "代理服务器运行时出错", e)
            } finally {
                closeServerSocket()
                Log.i(TAG, "代理服务器已停止")
            }
        }
    }

    private fun handleConnection(clientSocket: Socket) {
        connectionScope.launch {
            val requestId = System.currentTimeMillis()
            try {
                Log.d(TAG, "[$requestId] 客户端地址: ${clientSocket.inetAddress.hostAddress}:${clientSocket.port}")
                Log.d(TAG, "[$requestId] 本地地址: ${clientSocket.localAddress.hostAddress}:${clientSocket.localPort}")
                
                processConnection(clientSocket, requestId)
            } catch (e: Exception) {
                Log.e(TAG, "[$requestId] 处理连接时出错", e)
            } finally {
                closeSocket(clientSocket)
            }
        }
    }

    private fun processConnection(clientSocket: Socket, requestId: Long) {
        val input = getClientInput(clientSocket) ?: return
        val firstLine = readFirstLine(input) ?: return
        
        logRequest(requestId, firstLine)
        val headers = readHeaders(input, requestId)
        
        val parts = firstLine.split(" ")
        if (parts.size >= 3) {
            processRequest(parts, headers, input, clientSocket, requestId)
        } else {
            handleInvalidRequest(firstLine, clientSocket, requestId)
        }
    }

    private fun processRequest(
        parts: List<String>,
        headers: Map<String, String>,
        input: BufferedReader,
        clientSocket: Socket,
        requestId: Long
    ) {
        val (method, url, protocol) = parts
        Log.i(TAG, "[$requestId] 请求方法: $method, URL: $url, 协议: $protocol")

        try {
            when (method) {
                "CONNECT" -> handleHttpsConnection(clientSocket, url, requestId)
                else -> handleHttpConnection(clientSocket, parts.joinToString(" "), headers, input, requestId)
            }
        } catch (e: Exception) {
            val requestType = when(method) {
                "CONNECT" -> "HTTPS"
                else -> "HTTP"
            }
            Log.e(TAG, "[$requestId] 处理${requestType}请求失败", e)
        }
    }

    private fun handleInvalidRequest(firstLine: String, clientSocket: Socket, requestId: Long) {
        Log.e(TAG, "[$requestId] 无效的请求行: $firstLine")
        sendErrorResponse(clientSocket, "400 Bad Request", requestId)
    }

    private fun sendErrorResponse(clientSocket: Socket, error: String, requestId: Long) {
        try {
            clientSocket.getOutputStream().write(
                "HTTP/1.1 $error\r\n\r\n".toByteArray()
            )
        } catch (e: Exception) {
            Log.e(TAG, "[$requestId] 发送错误响应失败", e)
        }
    }

    private fun closeSocket(socket: Socket) {
        try {
            if (!socket.isClosed) {
                socket.close()
            }
        } catch (e: Exception) {
            Log.e(TAG, "关闭socket时出错", e)
        }
    }

    private fun getClientInput(clientSocket: Socket): BufferedReader? {
        try {
            return clientSocket.getInputStream().bufferedReader()
        } catch (e: Exception) {
            Log.e(TAG, "获取客户端输入流失败", e)
            return null
        }
    }

    private fun readFirstLine(input: BufferedReader): String? {
        try {
            val firstLine = input.readLine()
            if (firstLine == null) {
                Log.w(TAG, "请求首行为空")
                return null
            }
            return firstLine
        } catch (e: Exception) {
            Log.e(TAG, "读取请求首行失败", e)
            return null
        }
    }

    private fun logRequest(requestId: Long, firstLine: String) {
        val logPrefix = """
            ┌──────────── 请求 [$requestId] ────────────
            │ $firstLine
        """.trimIndent()
        Log.i(TAG, logPrefix)
    }

    private fun readHeaders(input: BufferedReader, requestId: Long): Map<String, String> {
        val headers = mutableMapOf<String, String>()
        try {
            var line: String?
            while (input.readLine().also { line = it } != null && line!!.isNotEmpty()) {
                Log.d(TAG, "│ $line")
                val index = line!!.indexOf(": ")
                if (index > 0) {
                    headers[line!!.substring(0, index)] = line!!.substring(index + 2)
                }
            }
        } catch (e: Exception) {
            Log.e(TAG, "读取请求头失败", e)
        }
        Log.d(TAG, "└────────────────────────────────────")
        return headers
    }

    private fun handleHttpConnection(
        clientSocket: Socket,
        firstLine: String,
        headers: Map<String, String>,
        clientInput: BufferedReader,
        requestId: Long
    ) {
        try {
            val parts = firstLine.split(" ")
            val url = parts[1]
            val hostAndPort = extractHostAndPort(url, 80)

            Log.i(TAG, """
                ┌──────────── HTTP 请求 [$requestId] ────────────
                │ 方法: ${parts[0]}
                │ URL: $url
                │ 协议: ${parts[2]}
                │ 目标主机: ${hostAndPort.first}:${hostAndPort.second}
                │ Headers:
                ${headers.entries.joinToString("\n") { "│   ${it.key}: ${it.value}" }}
                └────────────────────────────────────
            """.trimIndent())

            // 连接到目标服务器
            val targetSocket = try {
                Socket(hostAndPort.first, hostAndPort.second).apply {
                    soTimeout = 30000
                }
            } catch (e: Exception) {
                Log.e(TAG, "[$requestId] 连接目标服务器失败: ${hostAndPort.first}:${hostAndPort.second}", e)
                throw e
            }

            // 转发请求
            try {
                val targetOutput = targetSocket.getOutputStream()
                targetOutput.write("$firstLine\r\n".toByteArray())
                headers.forEach { (key, value) ->
                    targetOutput.write("$key: $value\r\n".toByteArray())
                }
                targetOutput.write("\r\n".toByteArray())

                // 如果有请求体，转发请求体
                val contentLength = headers["Content-Length"]?.toIntOrNull() ?: 0
                if (contentLength > 0) {
                    val buffer = ByteArray(contentLength)
                    var totalRead = 0
                    val inputStream = clientSocket.getInputStream()
                    while (totalRead < contentLength) {
                        val read = inputStream.read(buffer, totalRead, contentLength - totalRead)
                        if (read == -1) break
                        totalRead += read
                    }
                    Log.d(TAG, "[$requestId] 转发请求体: $totalRead bytes")
                    targetOutput.write(buffer, 0, totalRead)
                }
                targetOutput.flush()
            } catch (e: Exception) {
                Log.e(TAG, "[$requestId] 转发请求失败", e)
                throw e
            }

            // 读取并转发响应
            try {
                val targetInput = targetSocket.getInputStream().bufferedReader()
                val response = targetInput.readLine() ?: throw IOException("无响应")
                
                Log.i(TAG, """
                    ┌──────────── HTTP 响应 [$requestId] ────────────
                    │ 状态行: $response
                """.trimIndent())

                // 读取响应头
                val responseHeaders = mutableMapOf<String, String>()
                var headerLine: String? = null
                val headerBuilder = StringBuilder()
                while (targetInput.readLine().also { headerLine = it } != null && headerLine!!.isNotEmpty()) {
                    headerBuilder.append("│ $headerLine\n")
                    val index = headerLine!!.indexOf(": ")
                    if (index > 0) {
                        responseHeaders[headerLine!!.substring(0, index)] = headerLine!!.substring(index + 2)
                    }
                }
                Log.i(TAG, headerBuilder.toString() + "└────────────────────────────────────")

                // 转发响应给客户端
                val clientOutput = clientSocket.getOutputStream()
                clientOutput.write("$response\r\n".toByteArray())
                responseHeaders.forEach { (key, value) ->
                    clientOutput.write("$key: $value\r\n".toByteArray())
                }
                clientOutput.write("\r\n".toByteArray())

                // 转发响应体
                val buffer = ByteArray(BUFFER_SIZE)
                var bytesRead: Int
                var totalBytes = 0L
                val inputStream = targetSocket.getInputStream()
                while (inputStream.read(buffer).also { bytesRead = it } != -1) {
                    clientOutput.write(buffer, 0, bytesRead)
                    totalBytes += bytesRead
                }
                clientOutput.flush()
                Log.i(TAG, "[$requestId] 响应体转发完成: $totalBytes bytes")

            } catch (e: Exception) {
                Log.e(TAG, "[$requestId] 处理响应失败", e)
                throw e
            }

        } catch (e: Exception) {
            Log.e(TAG, "[$requestId] HTTP 连接处理失败", e)
            try {
                clientSocket.getOutputStream().write(
                    "HTTP/1.1 502 Bad Gateway\r\n\r\n".toByteArray()
                )
            } catch (e2: Exception) {
                Log.e(TAG, "[$requestId] 发送错误响应失败", e2)
            }
        }
    }

    private fun handleHttpsConnection(clientSocket: Socket, url: String, requestId: Long) {
        try {
            val hostAndPort = extractHostAndPort(url, 443)
            
            Log.i(TAG, """
                ┌──────────── HTTPS 请求 [$requestId] ────────────
                │ 主机: ${hostAndPort.first}:${hostAndPort.second}
                │ 类型: CONNECT
                └────────────────────────────────────
            """.trimIndent())

            // 生成服务器证书
            val certEntry = try {
                certificateManager.generateServerCertificate(hostAndPort.first)
            } catch (e: Exception) {
                Log.e(TAG, "[$requestId] 生成服务器证书失败", e)
                throw e
            }

            // 创建 SSL 上下文
            val sslContext = try {
                SSLContext.getInstance("TLS").apply {
                    val keyManagerFactory = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm())
                    val keyStore = KeyStore.getInstance(KeyStore.getDefaultType())
                    keyStore.load(null, null)
                    keyStore.setKeyEntry(
                        "cert",
                        certEntry.keyPair.private,
                        null,
                        arrayOf(certEntry.certificate)
                    )
                    keyManagerFactory.init(keyStore, null)
                    init(keyManagerFactory.keyManagers, null, null)
                }
            } catch (e: Exception) {
                Log.e(TAG, "[$requestId] 创建 SSL 上下文失败", e)
                throw e
            }

            // 发送连接成功响应
            try {
                clientSocket.getOutputStream().write(
                    "HTTP/1.1 200 Connection Established\r\n\r\n".toByteArray()
                )
                Log.i(TAG, "[$requestId] HTTPS 隧道建立成功")
            } catch (e: Exception) {
                Log.e(TAG, "[$requestId] 发送 HTTPS 连接确认失败", e)
                throw e
            }

            // 创建 SSL socket
            val sslSocket = try {
                sslContext.socketFactory.createSocket(
                    clientSocket,
                    hostAndPort.first,
                    hostAndPort.second,
                    true
                ) as SSLSocket
            } catch (e: Exception) {
                Log.e(TAG, "[$requestId] 创建服务器端 SSL socket 失败", e)
                throw e
            }

            // 开始 SSL 握手
            try {
                sslSocket.startHandshake()
                Log.i(TAG, "[$requestId] SSL 握手成功")
            } catch (e: Exception) {
                Log.e(TAG, "[$requestId] SSL 握手失败", e)
                throw e
            }

            // 创建到目标服务器的连接
            val targetSocket = try {
                SSLSocketFactory.getDefault().createSocket(
                    hostAndPort.first,
                    hostAndPort.second
                ) as SSLSocket
            } catch (e: Exception) {
                Log.e(TAG, "[$requestId] 连接目标 HTTPS 服务器失败", e)
                throw e
            }

            // 双向转发加密数据
            startBidirectionalForwarding(sslSocket, targetSocket, requestId)

        } catch (e: Exception) {
            Log.e(TAG, "[$requestId] HTTPS 连接处理失败", e)
            try {
                clientSocket.getOutputStream().write(
                    "HTTP/1.1 502 Bad Gateway\r\n\r\n".toByteArray()
                )
            } catch (e2: Exception) {
                Log.e(TAG, "[$requestId] 发送错误响应失败", e2)
            }
        }
    }

    private fun startBidirectionalForwarding(clientSocket: Socket, targetSocket: Socket, requestId: Long) {
        // 客户端 -> 服务器
        connectionScope.launch {
            try {
                forwardData(
                    clientSocket.getInputStream(),
                    targetSocket.getOutputStream(),
                    "客户端 -> 服务器",
                    requestId
                )
            } catch (e: Exception) {
                Log.e(TAG, "[$requestId] 客户端到服务器的转发失败", e)
            }
        }

        // 服务器 -> 客户端
        connectionScope.launch {
            try {
                forwardData(
                    targetSocket.getInputStream(),
                    clientSocket.getOutputStream(),
                    "服务器 -> 客户端",
                    requestId
                )
            } catch (e: Exception) {
                Log.e(TAG, "[$requestId] 服务器到客户端的转发失败", e)
            }
        }
    }

    private fun forwardData(
        input: InputStream,
        output: OutputStream,
        direction: String,
        requestId: Long
    ) {
        val buffer = ByteArray(BUFFER_SIZE)
        var bytesRead: Int
        var totalBytes = 0L
        try {
            while (input.read(buffer).also { bytesRead = it } != -1) {
                output.write(buffer, 0, bytesRead)
                output.flush()
                totalBytes += bytesRead
                if (totalBytes % (BUFFER_SIZE * 10) == 0L) {
                    Log.d(TAG, "[$requestId] $direction: 已转发 $totalBytes bytes")
                }
            }
            Log.i(TAG, "[$requestId] $direction: 转发完成，总共 $totalBytes bytes")
        } catch (e: Exception) {
            Log.e(TAG, "[$requestId] $direction: 数据转发失败", e)
            throw e
        }
    }

    private fun extractHostAndPort(url: String, defaultPort: Int): Pair<String, Int> {
        return when {
            url.startsWith("http://") || url.startsWith("https://") -> {
                val cleanUrl = url.substring(url.indexOf("//") + 2)
                val hostPart = cleanUrl.substringBefore('/')
                when {
                    hostPart.contains(":") -> Pair(hostPart.substringBefore(':'), hostPart.substringAfter(':').toInt())
                    else -> Pair(hostPart, defaultPort)
                }
            }
            url.contains(":") -> Pair(url.substringBefore(':'), url.substringAfter(':').toInt())
            else -> Pair(url, defaultPort)
        }
    }

    fun stop() {
        if (!running) {
            Log.w(TAG, "代理服务器已经停止")
            return
        }

        Log.i(TAG, "正在停止代理服务器...")
        running = false
        closeServerSocket()
        threadPool.shutdown()
        try {
            if (threadPool.awaitTermination(5, TimeUnit.SECONDS)) {
                Log.i(TAG, "代理服务器已成功停止")
            } else {
                Log.w(TAG, "代理服务器停止超时，强制关闭")
                threadPool.shutdownNow()
            }
        } catch (e: Exception) {
            Log.e(TAG, "停止代理服务器时出错", e)
            threadPool.shutdownNow()
        }
    }

    private fun closeServerSocket() {
        try {
            serverSocket?.close()
            serverSocket = null
        } catch (e: Exception) {
            Log.e(TAG, "关闭服务器socket时出错", e)
        }
    }

    private fun parsePacket(data: ByteArray): HttpRequest? {
        try {
            val input = String(data)
            val lines = input.split("\r\n")
            if (lines.isEmpty()) return null

            val firstLine = lines[0]
            val parts = firstLine.split(" ")
            if (parts.size < 3) return null

            val method = parts[0]
            val url = parts[1]
            val protocol = parts[2]

            // 解析请求头
            val headers = mutableMapOf<String, String>()
            var i = 1
            while (i < lines.size && lines[i].isNotEmpty()) {
                val index = lines[i].indexOf(": ")
                if (index > 0) {
                    headers[lines[i].substring(0, index)] = lines[i].substring(index + 2)
                }
                i++
            }

            // 解析主机和路径
            val hostAndPort = extractHostAndPort(url, when(method) {
                "CONNECT" -> 443
                else -> 80
            })
            
            return HttpRequest(
                method = method,
                url = url,
                host = hostAndPort.first,
                path = url.substringAfter(hostAndPort.first).ifEmpty { "/" },
                protocol = protocol,
                requestHeaders = headers,
                sourceIp = "127.0.0.1",
                sourcePort = port,
                destIp = hostAndPort.first,
                destPort = hostAndPort.second
            )
        } catch (e: Exception) {
            Log.e(TAG, "Error parsing packet", e)
            return null
        }
    }

    private fun getLogPrefix(method: String): String = when(method) {
        "CONNECT" -> "HTTPS"
        else -> "HTTP"
    }

    private fun handleClient(clientSocket: Socket) {
        try {
            Log.d(TAG, "收到新的客户端连接: ${clientSocket.inetAddress}")
            
            val input = clientSocket.getInputStream().bufferedReader()
            val firstLine = input.readLine()
            
            Log.d(TAG, "收到请求: $firstLine")
            
            // ... 处理请求的代码 ...
            
        } catch (e: Exception) {
            Log.e(TAG, "处理客户端请求失败", e)
        } finally {
            try {
                clientSocket.close()
            } catch (e: Exception) {
                Log.e(TAG, "关闭客户端连接失败", e)
            }
        }
    }

    private fun processRequest(request: HttpRequest, clientSocket: Socket) {
        Log.d(TAG, "开始处理请求: ${request.method} ${request.path}")
        // ... 其他代码保持不变 ...
    }
} 